#! coding: utf-8

__author__ = 'Anton Vakhrushev'

import wx


class ProgressError(Exception):
    pass


class ProgressStopIteration(ProgressError):
    pass


class ProgressSkipIteration(ProgressError):
    pass


class Progress(object):

    @staticmethod
    def cover(progress, *args, **kwargs):
        if progress:
            return progress
        return Progress(*args, **kwargs)

    def __init__(self, parent=None):
        self.dlg = wx.GenericProgressDialog(u'Внимание', u'Идет процесс...', 100, parent,
                                            wx.PD_APP_MODAL | wx.PD_CAN_ABORT | wx.PD_CAN_SKIP)
        self.dlg.SetSize((450, -1))
        self.dlg.Refresh()
        self.func = None
        self.iterable = None
        self.maximum = None
        self.message = None
        self.step = 0
        self.current = 0
        self.owner_count = 0

    def init(self, func, iterable=None, maximum=None, title=None, message=None):
        # check function
        if not callable(func):
            raise ProgressError(u'`func` must be callable')
        self.func = func
        # check count
        if maximum is None:
            try:
                self.maximum = len(iterable)
            except TypeError:
                self.maximum = None
        # check iterable
        self.iterable = iterable
        if self.iterable is None:
            self.maximum = None
        # check title
        if message is None:
            self.message = lambda index, item: unicode(item)
        elif isinstance(message, basestring):
            self.message = lambda index, item: unicode(message)
        elif callable(message):
            self.message = lambda index, item: unicode(message(index, item))
        else:
            raise ProgressError(u'{} is not appropriate type for `message`'.format(type(message)))
        self.current = 0
        self.step += 1
        self.dlg.SetTitle(u'[Шаг {}] {}'.format(self.step, title or u'Внимание'))
        self.dlg.Resume()

    def run(self):
        try:
            self.dlg.SetRange(self.maximum or 100)
            if self.iterable is not None:
                for item in self.iterable:
                    self.update(self.message(self.current, item))
                    yield self.func(item)
                    self.current += 1
            else:
                self.update(self.message(self.current, None))
                yield self.func()
                self.current += 1
        except ProgressSkipIteration:
            return

    def update(self, text):
        if len(text) > 50:
            text = u'...' + text[-50:]
        if self.maximum is None:
            text = u'[{}] {}'.format(self.current, text)
            self.dlg.Pulse(text)
        else:
            text = u'[{:.2%}] {}'.format(float(self.current) / self.maximum, text)
            self.dlg.Update(self.current, text)
        if self.dlg.WasCancelled():
            print 'stopped'
            raise ProgressStopIteration()
        if self.dlg.WasSkipped():
            print 'skipped'
            raise ProgressSkipIteration()

    def map(self, func, iterable, count=None, title=None, message=None):
        self.init(func, iterable, count, title, message)
        for item in self.run():
            yield item

    def wrap(self, iterable, count=None, title=None, message=None):
        self.init(lambda x: x, iterable, count, title, message)
        for item in self.run():
            yield item

    def do(self, func, title=None):
        self.init(func, None, None, title, title)
        return next(self.run())

    def close(self):
        self.dlg.Destroy()

    def __enter__(self):
        self.owner_count += 1

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.owner_count -= 1
        if not self.owner_count:
            self.close()